from fastapi import FastAPI, HTTPException
from fastapi.responses import FileResponse, HTMLResponse
from fastapi.staticfiles import StaticFiles
import os
import rsa
import tempfile
import re

app = FastAPI(title="RSA Key Manager")

# Mount static files directory
app.mount("/static", StaticFiles(directory="static"), name="static")

# Store temporary key paths
temp_files = {}

# 自定义函数用于修改密钥的头尾标识，去掉RSA
def modify_key_headers(key_content):
    # 将"BEGIN RSA"替换为"BEGIN"，将"END RSA"替换为"END"
    modified_content = re.sub(r'-----BEGIN RSA (PRIVATE|PUBLIC) KEY-----', r'-----BEGIN \1 KEY-----', key_content)
    modified_content = re.sub(r'-----END RSA (PRIVATE|PUBLIC) KEY-----', r'-----END \1 KEY-----', modified_content)
    return modified_content

@app.get("/", response_class=HTMLResponse)
async def root():
    with open("static/index.html", "r", encoding="utf-8") as f:
        return f.read()

@app.post("/generate")
async def generate_keys(key_size: int = 2048):
    try:
        # Generate new key pair
        (pubkey, privkey) = rsa.newkeys(key_size)
        
        # 获取原始密钥内容
        public_key_content = pubkey.save_pkcs1().decode('utf-8')
        private_key_content = privkey.save_pkcs1().decode('utf-8')
        
        # 修改密钥头尾标识
        modified_public_key = modify_key_headers(public_key_content)
        modified_private_key = modify_key_headers(private_key_content)
        
        # Create temporary files for keys
        pub_temp = tempfile.NamedTemporaryFile(delete=False, suffix='.pem', mode='wb')
        pub_temp.write(modified_public_key.encode('utf-8'))
        pub_temp.close()
        
        priv_temp = tempfile.NamedTemporaryFile(delete=False, suffix='.pem', mode='wb')
        priv_temp.write(modified_private_key.encode('utf-8'))
        priv_temp.close()
        
        # Store temp file paths for later use
        temp_files["public"] = pub_temp.name
        temp_files["private"] = priv_temp.name
        
        return {"message": "Keys generated successfully", 
                "public_key": modified_public_key,
                "private_key": modified_private_key}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Error generating keys: {str(e)}")

@app.get("/download/public")
async def download_public_key():
    if "public" not in temp_files or not os.path.exists(temp_files["public"]):
        raise HTTPException(status_code=404, detail="Public key not found. Generate keys first.")
    return FileResponse(temp_files["public"], filename="public_key.pem")

@app.get("/download/private")
async def download_private_key():
    if "private" not in temp_files or not os.path.exists(temp_files["private"]):
        raise HTTPException(status_code=404, detail="Private key not found. Generate keys first.")
    return FileResponse(temp_files["private"], filename="private_key.pem")

@app.on_event("shutdown")
def delete_temp_files():
    # Clean up temporary files when app shuts down
    for key, path in temp_files.items():
        if os.path.exists(path):
            os.unlink(path)

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)